%% Data
clear all

addpath('...\replication\Section 6\01_data');

% Commutes
L_ni = dlmread('...\replication\Section 6\01_data\commutes_ij_final.csv'); 
R_n = sum(L_ni,2);
L_i = sum(L_ni);

% Identify and correct
R_n_p = (R_n~=0);
L_ni = L_ni(R_n_p,R_n_p');

L_i_p = (L_i~=0);
L_ni = L_ni(L_i_p',L_i_p);

% Travel
d_ni = dlmread('...\replication\Section 6\01_data\trav_dist_ij_final.csv'); 
d_ni = d_ni(R_n_p,R_n_p');
d_ni = d_ni(L_i_p',L_i_p);

% Identify and correct
d_n = sum(d_ni,2);
d_n_p = (d_n~=0);

d_ni = d_ni(d_n_p,d_n_p');
L_ni = L_ni(d_n_p,d_n_p');

% Set vars after corrections
R_n = sum(L_ni,2);
L_i = sum(L_ni);

% Vars n
n_vars = dlmread('...\replication\Section 6\01_data\i_vars.csv'); 
n_vars = n_vars(R_n_p,1:end);
n_vars = n_vars(L_i_p',1:end);
n_vars = n_vars(d_n_p,1:end);

H_R = n_vars(:,5);
Q_n = n_vars(:,6).*(90.9*12);
v_n = n_vars(:,7);
%v_n=v_n./v_n(1:1);
xy_i = [n_vars(:,2),n_vars(:,3)];

% Shock
tau_ni = dlmread('...\replication\Section 6\01_data\tau_ij.csv'); 
tau_ni = tau_ni(R_n_p,R_n_p');
tau_ni = tau_ni(L_i_p',L_i_p);
tau_ni = tau_ni(d_n_p,d_n_p');

%% Parameters

L=sum(sum(L_ni),2);    % Total population in 2000 (benchmark) based on 330prefs+ROW
alpha=0.6 ;        % Consumption share (non-land share) in utility
Omega= 0.5;         % Location taste heterogeneity parameter
xi = 0.113;       % semi-elasticity of commutes to distance
phi = 0.43;            % semi-elasticity of trade to distance
sigma = 4;
epsilon = 3.3;

n=length(R_n);    % size of locations

%% Commutes and trade scalars

euc_dist = pdist2(xy_i,xy_i,'euclidean'); % Limit counterfactual to GCD
euc_dist = 100.*euc_dist; % Make km
k = find(d_ni<euc_dist);
d_ni(k)=euc_dist(k);
d_ni(d_ni<0.1) = 0.1;

d_ni = d_ni./90; % Travel speed on reg roads

% Kappa
kappa = (d_ni).^(-xi);
% Trade
zeta = (d_ni).^(-phi);

%% Workplace wages 
lambda_ni_i = L_ni./L_i;
w_i = sum(lambda_ni_i.*v_n);

%% Calibrate A

gaptotal = 1;

A_i_0=ones(1,n);

LHS = w_i.*L_i;
inc_RHS = v_n.*R_n;
fixed_comp = L_i.*(zeta.*w_i).^(1-sigma);

it = 1
while it < 1000

share = (fixed_comp.*A_i_0)./(sum(fixed_comp.*A_i_0,2));
RHS = sum(inc_RHS.*share);

% Stopping mechanism
gap = (LHS-RHS)';
gap=abs(gap);
gap=max(gap)

if gap <= gaptotal;
    
    it=10000;
    display('>>>> Productivity System Converged <<<<');
    A_i=A_i_0;
else
    
    if isnan(gap)==1;
    A_i_00=0.95+(1.05-0.95).*rand(length(LHS),1);
    A_i_0=0.95+(1.05-0.95).*rand(length(LHS),1);
    end;
    
    ratio = (LHS./RHS);
    
    if max(isnan(RHS))==1;
    stop
    end;
    
    A_i_00= ratio.*A_i_0;    
    A_i_0 = A_i_00;
    
        % Choose units for B_i such that geometric mean B_i=1;
 A_i_0=A_i_0./geomean(A_i_0);
  
end
it=it+1 
end
  
%% Calibrate B

gaptotal = 10.^-5;

B_n_0=ones(n,1);

LHS = R_n./L;

num = L_i.*(zeta.*w_i./A_i).^(1-sigma);
pi_ni = num./sum(num);
pi_nn = diag(pi_ni);

rat1 = (L_i'./pi_nn).^(-alpha*epsilon./(sigma-1));
rat2 = (R_n./H_R).^(-epsilon*(1-alpha));
A_n = A_i';
w_n = w_i';
fixed_comp = (kappa.^(-epsilon)) .* rat1 .* A_n.^(alpha*epsilon) .* w_n.^(-epsilon*(1-alpha)) .* rat2 .* w_i.^epsilon;

it = 1
while it < 1000

RHS = sum(fixed_comp.*B_n_0,2)./(sum(sum(fixed_comp.*B_n_0,2)));

% Stopping mechanism
gap =(LHS-RHS);
gap=abs(gap);
gap=max(gap)

if gap <= gaptotal;
    
    it=10000;
    display('>>>> Amenity System Converged <<<<');
    B_n=B_n_0;
else
    
    ratio = (LHS./RHS);
    
    if max(isnan(RHS))==1;
    stop
    end;
    
    B_n_00= ratio.*B_n_0;    
    B_n_0 = B_n_00; 
    
        % Choose units for B_i such that geometric mean B_i=1;
  B_n_0=B_n_0./geomean(B_n_0);
  
end
it=it+1 
end


% Export data
csvwrite('...\replication\Section 6\01_data\A_i.csv',A_i);
csvwrite('...\replication\Section 6\01_data\B_n.csv',B_n);
csvwrite('...\replication\Section 6\01_data\w_i.csv',w_i);